from __future__ import (absolute_import, division, print_function,unicode_literals)
import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])
import pandas as pd
import matplotlib.pyplot as plt
import backtrader as bt
from sqlalchemy import create_engine
from lib import *  # 导入lib中所有公开的类
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
from queue import Queue
import time

# 数据库连接信息
'''
db_config = {
    'host': '127.0.0.1',         # 数据库主机地址
    'user': 'root',     # 数据库用户名
    'password': 'root', # 数据库密码
    'database': 'stock_db', # 数据库名称
    'port': 3306                 # 数据库端口，默认3306
}
'''
db_config = {
    'host': '172.30.168.43',         # 数据库主机地址
    'user': 'root',     # 数据库用户名
    'password': 'DAzA8nuLfDCQfD2n', # 数据库密码
    'database': 'stock_db', # 数据库名称
    'port': 3306                 # 数据库端口，默认3306
}

def create_engine_for_thread():
    """为每个线程创建独立的数据库连接引擎"""
    return create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}")

def run_strategy_for_stock(ts_code, engine=None):
    """为单个股票运行策略 - 线程安全版本"""
    print(f"\n=== 开始处理股票: {ts_code} (线程: {threading.current_thread().name}) ===")
    
    # 为每个线程创建独立的数据库引擎
    if engine is None:
        engine = create_engine_for_thread()
    
    try:
        # Create a cerebro entity
        cerebro = bt.Cerebro()

        # Add a strategy with conservative parameters
        cerebro.addstrategy(LowPointStrategy, position_ratio=0.9)  # 使用90%资金比例

        # Create a Data Feed
        query = """
            SELECT 
                FROM_UNIXTIME(trade_date, %s) AS trade_date,
                open,
                high,
                low,
                close,
                vol AS volume
            FROM 
                base_daily 
            WHERE 
                ts_code = %s
                AND trade_date > UNIX_TIMESTAMP(DATE_SUB(NOW(3), INTERVAL 24 MONTH))
            ORDER BY 
                trade_date ASC
            """
        
        date_format = '%Y-%m-%d %H:%i:%s.%f'
        
        dataframe = pd.read_sql(query, engine, params=(date_format, ts_code), index_col='trade_date', parse_dates=['trade_date'])
        
        if dataframe.empty:
            print(f"股票 {ts_code} 没有数据，跳过")
            return None
            
        dataframe.index = pd.to_datetime(dataframe.index)
        
        # 确保列的顺序正确
        dataframe = dataframe[['open', 'high', 'low', 'close', 'volume']]
        
        # 适配Backtrader的数据格式
        data = bt.feeds.PandasData(
            dataname=dataframe,
            datetime=None,      # 使用索引作为时间
            open=0,             # open列索引
            high=1,             # high列索引
            low=2,              # low列索引
            close=3,            # close列索引
            volume=4,           # volume列索引
            openinterest=-1     # 无未平仓合约，设为-1
        )
        
        # Add the Data Feed to Cerebro
        cerebro.adddata(data, name=ts_code)

        # Set our desired cash start
        cerebro.broker.setcash(100000.0)
        cerebro.broker.setcommission(0.0003)

        # Print out the starting conditions
        start_value = cerebro.broker.getvalue()
        print(f'股票 {ts_code} - 初始资产: {start_value:.2f}')
        
        # Run over everything
        cerebro.run()

        # Print out the final result
        final_value = cerebro.broker.getvalue()
        profit = final_value - start_value
        profit_rate = (final_value / start_value - 1) * 100
        
        print(f'股票 {ts_code} - 最终资产: {final_value:.2f}')
        print(f'股票 {ts_code} - 盈亏: {profit:.2f} ({profit_rate:.2f}%)')
        
        return {
            'ts_code': ts_code,
            'start_value': start_value,
            'final_value': final_value,
            'profit': profit,
            'profit_rate': profit_rate
        }
        
    except Exception as e:
        print(f"处理股票 {ts_code} 时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        # 清理资源
        if 'engine' in locals():
            engine.dispose()

def process_stock_batch(stock_list, max_workers=4):
    """批量处理股票"""
    results = []
    
    # 使用线程池
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_stock = {
            executor.submit(run_strategy_for_stock, ts_code): ts_code 
            for ts_code in stock_list
        }
        
        # 收集结果
        for future in as_completed(future_to_stock):
            stock_code = future_to_stock[future]
            try:
                result = future.result()
                if result:
                    results.append(result)
            except Exception as exc:
                print(f'股票 {stock_code} 生成异常: {exc}')
    
    return results

if __name__ == '__main__':
    # 获取所有股票代码
    stock_sql = """
    SELECT ts_code 
    FROM base_stock_basic 
    WHERE (ts_code LIKE %s OR ts_code LIKE %s)
    """
    
    try:
        engine = create_engine_for_thread()
        stock_arr = pd.read_sql(stock_sql, engine, params=('%SZ%', '%SH%'))
        print("获取到的股票列表:")
        print(stock_arr)
        
        if stock_arr.empty:
            print("没有找到符合条件的股票")
            sys.exit(1)
        
        # 获取股票代码列表
        stock_codes = stock_arr['ts_code'].tolist()
        print(f"\n准备处理 {len(stock_codes)} 只股票")
        
        # 设置线程数（建议根据CPU核心数调整）
        max_workers = min(8, len(stock_codes))  # 最多8个线程，或股票数
        
        # 批量处理股票
        print(f"使用 {max_workers} 个线程开始处理...")
        start_time = time.time()
        
        results = process_stock_batch(stock_codes, max_workers=max_workers)
        
        end_time = time.time()
        print(f"\n处理完成，耗时: {end_time - start_time:.2f} 秒")
        
        # 打印汇总结果
        print("\n=== 策略执行汇总 ===")
        if results:
            df_results = pd.DataFrame(results)
            df_results = df_results.sort_values('profit_rate', ascending=False)
            print(df_results.to_string(index=False))
            
            # 计算总体表现
            total_start = sum([r['start_value'] for r in results])
            total_final = sum([r['final_value'] for r in results])
            total_profit = total_final - total_start
            avg_profit_rate = sum([r['profit_rate'] for r in results]) / len(results)
            
            print(f"\n总体表现:")
            print(f"总投资: {total_start:.2f}")
            print(f"总价值: {total_final:.2f}")
            print(f"总盈亏: {total_profit:.2f}")
            print(f"平均收益率: {avg_profit_rate:.2f}%")
            print(f"成功处理股票数: {len(results)}/{len(stock_codes)}")
        else:
            print("没有成功执行任何股票策略")
            
    except Exception as e:
        print(f"获取股票列表时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
